/**
 * @file discrete_sampler.h
 *
 * Sampler to draw samples from discrete distribution
 *
 * @author Dahua Lin
 */

#ifdef _MSC_VER
#pragma once
#endif

#ifndef GMLPP_DISCRETE_SAMPLER_H_
#define GMLPP_DISCRETE_SAMPLER_H_

#include <gmlpp/core/basic_rng.h>
#include <stdexcept>

namespace gml
{

	/********************************************
	 *
	 *  Simple sample (w/o replacement)
	 *
	 ********************************************/

	/**
	 * sampler for w/o replacement sampling
	 */
	template<typename TInt>
	class WORSampler : private noncopyable
	{
	public:
		WORSampler(const TInt& n)
		: m_n(n), m_idx(0), m_seq(new TInt[n])
		{
			for (TInt i = 0; i < n; ++i) m_seq[i] = i;
		}

		~WORSampler()
		{
			delete[] m_seq;
		}

	public:
		GML_ENSURE_INLINE void reset()
		{
			m_idx = 0;
		}

		GML_ENSURE_INLINE bool is_end() const
		{
			return m_idx == m_n;
		}

		GML_ENSURE_INLINE TInt n_remain() const
		{
			return m_n - m_idx;
		}

		template<class Rng>
		TInt next(Rng& rg)
		{
			TInt i = (TInt)rg.randi((uint32_t)(m_n - m_idx));
			swap_val(m_idx, m_idx + i);
			return m_seq[m_idx++];
		}

		template<class Rng>
		TInt safe_next(Rng& rg)
		{
			if (is_end())
				throw std::logic_error("The wor_sampler is ended.");

			return next(rg);
		}

	private:
		void swap_val(TInt i, TInt j)
		{
			TInt tv = m_seq[i];
			m_seq[i] = m_seq[j];
			m_seq[j] = tv;
		}

	private:
		TInt m_n;
		TInt m_idx;
		TInt* m_seq;

	}; // end class WORSampler


	/**
	 * Draw two integers from [0, n) without replacement
	 */
	template<class Rng, typename TInt>
	inline void randpick_pair(Rng& rg, TInt n, TInt& v1, TInt& v2)
	{
		v1 = rg.randi(n);
		v2 = rg.randi(n-1);

		if (v2 >= v1) ++ v2;
	}


	/********************************************
	 *
	 *  Draw from discrete distribution
	 *
	 ********************************************/


	template<typename TInt>
	TInt ddsample(double u, TInt K, const double *p)
	{
		TInt i = 1;
		double v = *p;

		while (u > v && i < K) v += p[i++];

		return i - 1;
	}

	template<class Rng, typename TInt>
	GML_ENSURE_INLINE inline TInt ddsample(Rng& rg, TInt K, const double *p)
	{
		return ddsample(rg.randu(), K, p);
	}

	template<typename TInt>
	TInt ddsample_w(double u, TInt K, const double *w)
	{
		double sw = 0;
		for (TInt k = 0; k < K; ++k) sw += w[k];

		TInt i = 1;
		double v = *w;
		double uw = u * sw;

		while (uw > v && i < K) v += w[i++];

		return i - 1;
	}

	template<class Rng, typename TInt>
	GML_ENSURE_INLINE inline TInt ddsample_w(Rng& rg, TInt K, const double *w)
	{
		return ddsample_w(rg.randu(), K, w);
	}



}

#endif /* DISCRETE_SAMPLER_H_ */
